import org.apache.atlas.AtlasClientV2;
import org.apache.atlas.AtlasServiceException;
import org.apache.atlas.model.typedef.*;
import org.apache.atlas.type.AtlasTypeUtil;

import java.util.Arrays;
import java.util.Collections;

/**
 * Created by sqy on 2018/12/12.
 */
public class TypeOperation {
     private final AtlasClientV2 atlasClientV2;
     private static final String PURPOSE_TYPE = "purpose";
     private static final String TASK_TYPE = "task";
     private static final String USEDDATA_TYPE = "usedData";

    public TypeOperation(AtlasClientV2 atlasClientV2) {
        this.atlasClientV2 = atlasClientV2;
    }

    public void createTypes() throws AtlasServiceException {
        AtlasTypesDef atlasTypesDef = createTypeDefintions();
        atlasClientV2.createAtlasTypeDefs(atlasTypesDef);
    }

    private AtlasTypesDef createTypeDefintions() {
        AtlasEntityDef purposeType = AtlasTypeUtil.createClassTypeDef(PURPOSE_TYPE,PURPOSE_TYPE,"1.0",null,
                                     AtlasTypeUtil.createUniqueRequiredAttrDef("name","string"),
                                     AtlasTypeUtil.createOptionalAttrDef("name","string"),
                                     AtlasTypeUtil.createOptionalAttrDef("description","string"),
                                     AtlasTypeUtil.createOptionalAttrDef("createTime","Date"),
                                     AtlasTypeUtil.createOptionalAttrDef("createBy","String"));

        AtlasEntityDef taskType = AtlasTypeUtil.createClassTypeDef(TASK_TYPE,TASK_TYPE,"1.0",null,
                                  AtlasTypeUtil.createUniqueRequiredAttrDef("name","string"),
                                  AtlasTypeUtil.createOptionalAttrDef("name","string"),
                                  AtlasTypeUtil.createOptionalAttrDef("description","string"),
                                  AtlasTypeUtil.createOptionalAttrDef("createTime","Date"),
                                  AtlasTypeUtil.createOptionalAttrDef("createBy","String"));

        AtlasEntityDef uesdDataType = AtlasTypeUtil.createClassTypeDef(USEDDATA_TYPE,USEDDATA_TYPE,"1.0",null,
                AtlasTypeUtil.createUniqueRequiredAttrDef("name","string"),
                AtlasTypeUtil.createOptionalAttrDef("name","string"),
                AtlasTypeUtil.createOptionalListAttrDef("inputs", AtlasBaseTypeDef.getArrayTypeName(PURPOSE_TYPE)),
                AtlasTypeUtil.createOptionalListAttrDef("outputs",AtlasBaseTypeDef.getArrayTypeName(TASK_TYPE)),
                AtlasTypeUtil.createOptionalAttrDef("description","string"),
                AtlasTypeUtil.createOptionalAttrDef("createTime","Date"),
                AtlasTypeUtil.createOptionalAttrDef("createBy","String"));

        return AtlasTypeUtil.getTypesDef(Collections.emptyList(),
                Collections.emptyList(),
                Collections.emptyList(),
                Arrays.asList(purposeType, uesdDataType, taskType));
    }
}
